# -*- coding: utf-8 -*-
# @Author  : LG
# 使用了与torchvision中VGG相同的搭建方式.

from torch import nn
import os
import torch

# 参数配置,标准的darknet19参数.
cfg = [32, 'M', 64, 'M', 128, 64, 128, 'M', 256, 128, 256, 'M',
       512, 256, 512, 256, 512, 'M', 1024, 512, 1024, 512, 1024]

def make_layers(cfg, in_channels=3, batch_norm=True):
    layers = []
    flag = True             # 用于变换卷积核大小,(True选后面的,False选前面的)
    in_channels= in_channels
    for v in cfg:
        if v == 'M':
            layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
        else:
            layers.append(nn.Conv2d(in_channels = in_channels,
                                   out_channels= v,
                                   kernel_size=(1, 3)[flag],
                                   stride=1,
                                   padding=(0,1)[flag],
                                   bias=False))
            if batch_norm:
                layers.append(nn.BatchNorm2d(v))
            in_channels = v

            layers.append(nn.LeakyReLU(negative_slope=0.1, inplace=True))
        flag = not flag
    return nn.Sequential(*layers)

class Darknet19(nn.Module):
    def __init__(self, num_classes=1000, in_channels=3, batch_norm=True):
        super(Darknet19, self).__init__()
        self.features = make_layers(cfg, in_channels=in_channels, batch_norm=batch_norm)
        self.classifier = nn.Sequential(
            nn.Conv2d(1024, num_classes, kernel_size=1, stride=1),
            nn.AdaptiveAvgPool2d(output_size=(1)),
            nn.Softmax(dim=0)
        )
        # 导入预训练模型
        self.load_weight()

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        x = torch.squeeze(x)
        return x

    def load_weight(self):
        weight_file = 'weights/darknet19-deepBakSu-e1b3ec1e.pth'
        dic = {}
        for now_keys, values in zip(self.state_dict().keys(), torch.load(weight_file).values()):
            dic[now_keys]=values
        self.load_state_dict(dic)

if __name__ == '__main__':
    net = Darknet19(num_classes=1000)
    print(net)
    x = torch.zeros((2,3,224,224))
    out = net(x)
    print(out.size())
